clear; clc; close all;
angstrom = 1e-10;  %[m]
a0 = 5.29177e-11;  %[m]
cmn1 = 0.0299792458e12; %[Hz/cm^-1]
foffset = -0.05;     %[GHz] wavemeter offset(-0.16 GHz) + AOM(+0.11GHz) 

fnames = dir('data/2020*K2*.mat');
Nfile = length(fnames);
dataAll = cell(1, Nfile);
disp('---------------------------------------------');
disp(['There are ',num2str(Nfile), ' files in total']);
dataxAll = [];
dataK2All = [];
dataK2_bkg_All = [];
dataNcyc = [];

for i = 1:Nfile
    filename4load = [fnames(i).folder '/' fnames(i).name];
    load(filename4load);
%     disp(['=> The ', num2str(i), 'th file Loaded.']);
    dataxi = xVarList1+foffset;                     %[GHz] 
    dataK2i = xVarCounts_K_2;             %ion counts of Rb2
    dataK2i_bkg = xVarCountsBkgd_K_2;     %bkgnd ion counts of Rb2
    dataNcyci = cycleNumxVar;               %cycle number
    dataxyi = [dataxi dataK2i dataK2i_bkg dataNcyci];
    dataAll{i} = dataxyi; 
    
    dataxAll = [dataxAll dataxi'];
    dataK2All = [dataK2All dataK2i'];
    dataK2_bkg_All = [dataK2_bkg_All dataK2i_bkg'];
    dataNcyc = [dataNcyc dataNcyci'];
end
disp('---------------------------------------------');


%%%------plot-------
f0 = 462800;
fcorrection = -0.02;         %[GHz]
df = 0.4;                      %[GHz]
markersize = 6;
yplotmin = -0.5;
yplotmax = 17;

% h1 = figure('units','normalized','WindowStyle','docked');
Xplot = (dataxAll-f0);
% % Yplot = (dataRb2All-dataRb2_bkg_All)./dataNcyc;
% % Yplot_err = sqrt(dataRb2All+dataRb2_bkg_All)./dataNcyc;
Yplot = (dataK2All-dataK2_bkg_All)./dataNcyc;
Yplot_err = sqrt(dataK2All+dataK2_bkg_All)./dataNcyc;
% errorbar(Xplot, Yplot, Yplot_err, 'ok');
% % ylim([yplotmin, yplotmax]);
% xlabel(['f - ', num2str(f0),' (GHz)']);
% ylabel('Ion count per cycle');
% title('All Rb_2 REMPI data with the rotation transitions of B^1\Pi_u(v''=6) - X^1\Sigma_g(v"=0,J")');

f_Q_branch = fcorrection+[...
        462895.30 462901.05 462906.35 462911.13 462915.44 462919.28 ...
        462922.64 462925.52 462927.92 462929.84 462931.28 462932.23...
        462935.50]; %[GHz]
col = 8;                        %plot column
NplotQ = length(f_Q_branch);
row = ceil(NplotQ/col);          %plot row

N_list = [];
A_list = [];
A_list_err = [];
dx_list = [];
dx_list_err = [];
xdataAll = [];
ydataAll = [];
yerr_dataAll = [];

h1 = figure('position', [0, 500, 1000, 600]);
subplot(2, 1, 1);
for j = 1:5
    row = ceil(NplotQ/col);          %plot row
%     subplot(2, 4, j);
    fQc = f_Q_branch(j);
    dfmin = df;
    dfmax = df;
    if j==3
        dfmin = 0.35;
    end
    xfitdata = Xplot(Xplot<=fQc-f0+dfmax&Xplot>=fQc-f0-dfmin);
    yfitdata = Yplot(Xplot<=fQc-f0+dfmax&Xplot>=fQc-f0-dfmin);
    yfitdata_err = Yplot_err(Xplot<=fQc-f0+dfmax&Xplot>=fQc-f0-dfmin);

    %---fit data to a model----
    xc = fQc-f0;
    plot([xc xc],[-1 20],'--b','LineWidth', 0.5);
    hold on
    Aguess = max(yfitdata);
    dxguess = 0.03;
    if mod(13-j,2) ==1      
        fo = fitoptions('Method','NonlinearLeastSquares',...
            'Lower',[0, 0],...
            'Upper',[50, df],...
            'StartPoint',[Aguess, dxguess]);
        ft = fittype('A.*exp(-(x-xc).^2/2/dx^2)','problem','xc','options',fo);
        [curvex,gofx] = fit(xfitdata', yfitdata', ft,'problem', xc);
    else
        fo = fitoptions('Method','NonlinearLeastSquares',...
            'Lower',[0, 0, xc-0.10],...
            'Upper',[50, df, xc+0.10],...
            'StartPoint',[Aguess, xc, dxguess]);
        ft = fittype('A.*exp(-(x-xc).^2/2/dx^2)','options',fo);
        [curvex,gofx] = fit(xfitdata', yfitdata', ft);
    end
    N_list = [N_list 13-j];
    A_list = [A_list curvex.A];
    ci = confint(curvex);
    A_list_err = [A_list_err (ci(2,1)-ci(1,1))/4];
    dx_list = [dx_list curvex.dx];
    dx_list_err = [dx_list_err (ci(2,2)-ci(1,2))/4];
    curvex1 = xc-0.6:0.001:xc+0.6;
    plot(curvex1, curvex(curvex1), 'k', 'LineWidth', 1);
    %     xplot = fQc-f0;
    %     yplot = ones(size(xplot)).*18;
    %     stem(xplot, yplot ,'ok');
    %     hold off
    legend off
    hold on
    if mod(13-j,2) ==0
        errorbar(xfitdata, yfitdata, yfitdata_err, 'ok', 'MarkerSize', markersize,...
            'MarkerFaceColor', 'r', 'CapSize', 0, 'LineWidth', 1);
    else
        errorbar(xfitdata, yfitdata, yfitdata_err, 'ok', 'MarkerSize', markersize,...
            'CapSize', 0, 'LineWidth', 1);
    end
    xdataAll = [xdataAll xfitdata];
    ydataAll = [ydataAll yfitdata];
    yerr_dataAll=[yerr_dataAll yfitdata_err];
    ax = gca;
    ax.FontSize = 13;
    xlabel(['f - ', num2str(f0),' (GHz)']);
    ylabel('Ion count per cycle');
    xlim([94, 117]);
    xticks(94:2:117);
    ylim([yplotmin, yplotmax]);
end

% h2 = figure('position', [100, 500, 1100, 700]);
subplot(2, 1, 2);
for j = 6:NplotQ
    disp(' ');
    disp(['j = ', num2str(13-j)]);
    row = ceil(NplotQ/col);          %plot row
    %     subplot(2, 4, j-(NplotQ-4)+1);
    fQc = f_Q_branch(j);
    xfitdata = Xplot(Xplot<=fQc-f0+df&Xplot>=fQc-f0-df);
    yfitdata = Yplot(Xplot<=fQc-f0+df&Xplot>=fQc-f0-df);
    yfitdata_err = Yplot_err(Xplot<=fQc-f0+df&Xplot>=fQc-f0-df);
    
    if 13-j==6
        xfitdata6 = xfitdata;
        yfitdata6 = yfitdata;
        yfitdata6_err = yfitdata_err;
    end
    if 13-j==5
        yfitdata = yfitdata-3.8*exp(-(xfitdata-125.68).^2/2/0.07078^2);
    end
    
    %---fit data to a model----
    xc = fQc-f0;
    plot([xc xc],[-1 20],'--b','LineWidth', 0.5);
    hold on
    Aguess = max(yfitdata);
    dxguess = 0.03;
    if mod(13-j,2) ==1
        fo = fitoptions('Method','NonlinearLeastSquares',...
            'Lower',[0, 0],...
            'Upper',[50, df],...
            'StartPoint',[Aguess, dxguess]);
        ft = fittype('A.*exp(-(x-xc).^2/2/dx^2)','problem','xc','options',fo);
        [curvex,gofx] = fit(xfitdata', yfitdata', ft,'problem', xc)
    else
        fo = fitoptions('Method','NonlinearLeastSquares',...
            'Lower',[0, 0, xc-0.15],...
            'Upper',[50, df, xc+0.15],...
            'StartPoint',[Aguess, xc, dxguess]);
        ft = fittype('A.*exp(-(x-xc).^2/2/dx^2)','options',fo);
        [curvex,gofx] = fit(xfitdata', yfitdata', ft)
    end
    N_list = [N_list 13-j];
    A_list = [A_list curvex.A];
    ci = confint(curvex);
    A_list_err = [A_list_err (ci(2,1)-ci(1,1))/4];
    dx_list = [dx_list curvex.dx];
    dx_list_err = [dx_list_err (ci(2,2)-ci(1,2))/4];
    
    curvex1 = xc-0.6:0.001:xc+0.6;
    plot(curvex1, curvex(curvex1), 'k', 'LineWidth', 1);
    if 13-j==6
        curvex6 = curvex1;
        curvey6 = curvex(curvex6);
    end
    %     xplot = fQc-f0;
    %     yplot = ones(size(xplot)).*18;
    %     stem(xplot, yplot ,'ok');
    %     hold off
    hold on
    legend off
    if mod(13-j,2) ==0
        errorbar(xfitdata, yfitdata, yfitdata_err, 'ok', 'MarkerSize', markersize,...
            'MarkerFaceColor', 'r', 'CapSize', 0, 'LineWidth', 1);
    else
        errorbar(xfitdata, yfitdata, yfitdata_err, 'ok', 'MarkerSize', markersize,...
            'CapSize', 0, 'LineWidth', 1);
    end
    xdataAll = [xdataAll xfitdata];
    ydataAll = [ydataAll yfitdata];
    yerr_dataAll=[yerr_dataAll yfitdata_err];
    ax = gca;
    ax.FontSize = 13;
    xlabel(['f - ', num2str(f0),' (GHz)']);
    ylabel('Ion count per cycle');
    xlim([117, 140]);
    ylim([yplotmin, yplotmax]);
    
    
    %     figure(h1)
    %     subplot(2, 1, 2);
end

if 0        %enable it for showing zoom in inset plot  
    axes('Position',[0.21 0.7 0.15 0.2]);
    box on;
%     xc = f_Q_branch(7)-f0;
%     plot([xc xc],[-1 20],'--b','LineWidth', 1);
%     hold on
    plot(curvex6, curvey6, 'k', 'LineWidth', 1)
    hold on
    errorbar(xfitdata6, yfitdata6, yfitdata6_err, 'ok', 'MarkerSize', markersize+2,...
        'MarkerFaceColor', [102 8 32]/(102+8+32)*0.7, 'CapSize', 0, 'LineWidth', 1);
    hold off
    xlim([122.3, 122.9]);
    ylim([yplotmin, yplotmax]);
    ax = gca;
    ax.FontSize = 11;5
end

h2 = figure('position', [0, 500, 500, 250]);
odd = mod(N_list,2);
even = ~odd;
bar_x = N_list;
bar_y = (A_list).*dx_list*sqrt(2*pi)/0.03;
bar_y_err = bar_y.*min(sqrt((A_list_err./A_list).^2),1);%+(dx_list_err./dx_list).^2
bar(bar_x,bar_y, 'FaceColor', [47 108 167]/(47+108+167)*1.3, 'EdgeColor', [47 108 167]/(47+108+167)*1.);
% legend('Even', 'Odd');
yticks([0:70/5:70]);
yticklabels({'0','0.2','0.4', '0.6','0.8','1'})
hold on
bar_y1 = bar_y;
bar_y1(odd&bar_x<20) = 0;
bar(bar_x,bar_y1, 'FaceColor', [102 8 32]/(102+8+32)*0.7, 'EdgeColor', [102 8 32]/(102+8+32)*0.5);
errorbar(bar_x, bar_y, bar_y_err, '.', 'MarkerSize', 0.5,...
         'CapSize', 3, 'LineWidth', 1, 'MarkerEdgeColor',[1 1 1]*0.,...
         'Color',[1 1 1]*0.);
hold off
% xlabel('Rotational state N');
% ylabel('Integrated ion count per cycle');
% title('K_2 REMPI');


figure(h2)


set(h1,'Units','Inches');
pos = get(h1,'Position');
set(h1,'PaperPositionMode','Auto','PaperUnits','Inches','PaperSize',[pos(3), pos(4)])
print(h1,'plots\K2REMPI_spectrum','-dpdf','-r0')